import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error
from skopt import gp_minimize
from skopt.space import Real
from typing import Tuple, List, Dict
from deap import base, creator, tools, algorithms
import random

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

# ---------------------------
# 增强版数据预处理
# ---------------------------
def preprocess_data(df: pd.DataFrame) -> pd.DataFrame:
    """处理原始数据，提取男胎样本并计算达标时间"""
    required_columns = ['孕妇代码', 'Y染色体浓度', '检测孕周']
    if not all(col in df.columns for col in required_columns):
        raise ValueError("输入数据缺少必要的列：'孕妇代码', 'Y染色体浓度', '检测孕周'")
    
    def convert_gestational_week(week_str):
        try:
            if isinstance(week_str, (int, float)):
                return float(week_str)
            if '+' in week_str:
                w, d = week_str.lower().replace('w', '').split('+')
                return float(w) + float(d)/7
            return float(week_str.lower().replace('w', ''))
        except:
            return np.nan
    
    df['检测孕周'] = df['检测孕周'].apply(convert_gestational_week)
    df = df.dropna(subset=['检测孕周'])
    
    # 筛选男胎数据（Y染色体浓度≥4%为男胎）
    male_data = df[df['Y染色体浓度'] >= 0.04].copy()
    
    # 计算达标孕周
    male_data['达标孕周'] = male_data.groupby('孕妇代码')['检测孕周'].transform('min')
    male_data = male_data.dropna(subset=['达标孕周'])
    
    return male_data

# ---------------------------
# 增强版预测模型构建
# ---------------------------
def build_time_prediction_model(data: pd.DataFrame) -> Tuple[RandomForestRegressor, PolynomialFeatures, StandardScaler]:
    """建立达标时间预测模型"""
    features = ['年龄', '身高', '体重', '孕妇BMI', '怀孕次数', '生产次数']
    missing_features = [f for f in features if f not in data.columns]
    if missing_features:
        raise ValueError(f"输入数据缺少必要的特征列：{missing_features}")
    
    X = data[features]
    y = data['达标孕周']
    
    # 添加多项式特征和标准化
    poly = PolynomialFeatures(degree=2, include_bias=False)
    scaler = StandardScaler()
    X_poly = scaler.fit_transform(poly.fit_transform(X))
    
    # 训练随机森林模型
    X_train, X_test, y_train, y_test = train_test_split(
        X_poly, y, test_size=0.2, random_state=42
    )
    
    model = RandomForestRegressor(
        n_estimators=150,
        max_depth=12,
        min_samples_split=5,
        random_state=42,
        n_jobs=-1
    )
    model.fit(X_train, y_train)
    
    # 模型评估
    train_pred = model.predict(X_train)
    test_pred = model.predict(X_test)
    
    print("\n=== 模型评估 ===")
    print(f"训练集R²: {model.score(X_train, y_train):.3f}")
    print(f"测试集R²: {model.score(X_test, y_test):.3f}")
    print(f"训练集MAE: {mean_absolute_error(y_train, train_pred):.3f} 周")
    print(f"测试集MAE: {mean_absolute_error(y_test, test_pred):.3f} 周")
    print(f"训练集RMSE: {np.sqrt(mean_squared_error(y_train, train_pred)):.3f} 周")
    print(f"测试集RMSE: {np.sqrt(mean_squared_error(y_test, test_pred)):.3f} 周")
    
    return model, poly, scaler

# ---------------------------
# 遗传算法实现BMI分组优化
# ---------------------------
def optimize_bmi_groups_ga(data: pd.DataFrame, max_groups: int = 4, min_samples: int = 30, 
                          pop_size: int = 50, ngen: int = 30, cxpb: float = 0.7, mutpb: float = 0.2) -> List[float]:
    """使用遗传算法优化BMI分组"""
    # 重置creator以防重复定义
    if hasattr(creator, 'FitnessMin'):
        del creator.FitnessMin
    if hasattr(creator, 'Individual'):
        del creator.Individual
    
    data = data.sort_values('孕妇BMI').reset_index(drop=True)
    n = len(data)
    bmi_values = data['孕妇BMI'].values
    
    # 数据量校验
    if n < min_samples * max_groups:
        raise ValueError(f"数据量不足，至少需要{min_samples*max_groups}个样本")
    
    # 创建遗传算法框架
    creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
    creator.create("Individual", list, fitness=creator.FitnessMin)
    
    toolbox = base.Toolbox()
    
    # 定义个体生成函数
    def generate_individual():
        # 生成随机分割点 (1到max_groups-1个分割点)
        num_splits = random.randint(1, max_groups-1)
        splits = sorted([random.uniform(bmi_values.min(), bmi_values.max()) for _ in range(num_splits)])
        # 确保分割点不重复且数量正确
        splits = list(sorted(set(splits)))
        # 确保分割点数量在有效范围内
        if len(splits) < 1:
            splits = [random.uniform(bmi_values.min(), bmi_values.max())]
        elif len(splits) > max_groups-1:
            splits = splits[:max_groups-1]
        return creator.Individual(splits)
    
    # 注册遗传算法操作
    toolbox.register("individual", generate_individual)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)
    
    # 定义交叉和变异操作
    def cxTwoPointCopy(ind1, ind2):
        """两点交叉操作"""
        size1 = len(ind1)
        size2 = len(ind2)
        
        # 如果个体长度小于2，直接返回父代
        if size1 < 2 or size2 < 2:
            return ind1, ind2
        
        # 确保交叉点在有效范围内
        size = min(size1, size2)
        cxpoint1 = random.randint(1, size-1)
        cxpoint2 = random.randint(1, size-1)
        
        # 确保cxpoint1 < cxpoint2
        if cxpoint2 < cxpoint1:
            cxpoint1, cxpoint2 = cxpoint2, cxpoint1
        elif cxpoint1 == cxpoint2:
            # 如果两个点相同，调整cxpoint2
            if cxpoint2 < size-1:
                cxpoint2 += 1
            else:
                cxpoint1 -= 1
        
        # 执行交叉
        ind1[cxpoint1:cxpoint2], ind2[cxpoint1:cxpoint2] = ind2[cxpoint1:cxpoint2].copy(), ind1[cxpoint1:cxpoint2].copy()
        
        return ind1, ind2
    
    def mutUniform(individual, bmi_values, indpb):
        """均匀变异操作（使用改进的变异函数）"""
        for i in range(len(individual)):
            if random.random() < indpb:
                # 在当前分割点附近进行变异，变异幅度基于BMI范围
                mutation_range = (bmi_values.max() - bmi_values.min()) * 0.1
                individual[i] += random.uniform(-mutation_range, mutation_range)
                # 确保变异后的值在合理范围内
                individual[i] = max(bmi_values.min(), min(bmi_values.max(), individual[i]))
        
        # 确保分割点有序且不重复
        individual = list(sorted(set(individual)))
        return individual,
    
    toolbox.register("mate", cxTwoPointCopy)
    toolbox.register("mutate", mutUniform, bmi_values=bmi_values, indpb=0.2)  # 使用改进的变异函数
    toolbox.register("select", tools.selTournament, tournsize=3)
    toolbox.register("evaluate", evaluate_individual, data=data, bmi_values=bmi_values, 
                    max_groups=max_groups, min_samples=min_samples)
    
    # 初始化种群
    pop = toolbox.population(n=pop_size)
    
    # 评估初始种群
    fitnesses = list(map(toolbox.evaluate, pop))
    for ind, fit in zip(pop, fitnesses):
        ind.fitness.values = fit
    
    # 运行遗传算法
    algorithms.eaSimple(pop, toolbox, cxpb=cxpb, mutpb=mutpb, ngen=ngen, verbose=False)
    
    # 获取最佳个体
    best_individual = tools.selBest(pop, k=1)[0]
    bounds = sorted([bmi_values.min()] + best_individual + [bmi_values.max()])
    
    return bounds

def evaluate_individual(individual, data, bmi_values, max_groups, min_samples):
    """评估个体的适应度（修复风险计算逻辑）"""
    # 添加最小和最大边界
    bounds = [bmi_values.min()] + sorted(individual) + [bmi_values.max()]
    
    # 检查分组数是否超过最大值
    if len(bounds) - 1 > max_groups:
        return float('inf'),
    
    # 计算每个组的风险（使用优化后的最佳时间）
    total_risk = 0
    for i in range(len(bounds)-1):
        lower, upper = bounds[i], bounds[i+1]
        group_data = data[(data['孕妇BMI'] >= lower) & (data['孕妇BMI'] < upper)]
        
        # 惩罚样本量不足的组
        if len(group_data) < min_samples:
            return float('inf'),
        
        # 使用优化得到的最佳时间计算风险
        group_risk = calculate_group_risk(group_data)
        total_risk += group_risk
    
    return total_risk,

# ---------------------------
# 其他辅助函数
# ---------------------------
def calculate_group_risk(group_data: pd.DataFrame) -> float:
    """计算单组的风险值（使用高斯过程优化）"""
    def time_objective(t: float) -> float:
        # 确保t是标量值
        t_value = t[0] if isinstance(t, (list, np.ndarray)) else t
        pred_prob = predict_detection_probability(group_data, t_value)
        risk_weight = 1 if t_value <= 12 else (2 if t_value <= 27 else 5)
        return (1 - pred_prob) * risk_weight  # 最小化风险
    
    # 高斯过程优化
    res = gp_minimize(
        time_objective,
        [Real(8, 40)],
        n_calls=50,
        random_state=42,
        acq_func='EI',
        verbose=False
    )
    return res.fun

def predict_detection_probability(group_data: pd.DataFrame, t: float) -> float:
    """预测在时间t的达标概率"""
    t_value = t[0] if isinstance(t, (list, np.ndarray)) else t
    early_records = group_data[group_data['检测孕周'] <= t_value]
    return (early_records['Y染色体浓度'] >= 0.04).mean() if len(early_records) > 0 else 0.0

# ---------------------------
# 修改后的主流程
# ---------------------------
def main(df: pd.DataFrame) -> pd.DataFrame:
    """主优化流程"""
    print("\n=== 数据预处理 ===")
    male_data = preprocess_data(df)
    print(f"有效男胎样本数: {len(male_data)}")
    
    print("\n=== 模型构建 ===")
    time_model, poly, scaler = build_time_prediction_model(male_data)
    
    print("\n=== BMI分组优化 (遗传算法) ===")
    bmi_bounds = optimize_bmi_groups_ga(male_data)
    print(f"获得BMI分组边界: {bmi_bounds}")
    
    # 为每个分组确定最佳检测时间
    groups = []
    for i in range(len(bmi_bounds)-1):
        lower, upper = bmi_bounds[i], bmi_bounds[i+1]
        group_data = male_data[
            (male_data['孕妇BMI'] >= lower) & 
            (male_data['孕妇BMI'] < upper)
        ]
        
        if len(group_data) < 30:  # 忽略小样本组
            continue
            
        # 使用高斯过程优化找到最佳检测时间
        def objective(t):
            t_value = t[0] if isinstance(t, (list, np.ndarray)) else t
            prob = predict_detection_probability(group_data, t_value)
            risk_weight = 1 if t_value <= 12 else (2 if t_value <= 27 else 5)
            return (1 - prob) * risk_weight
        
        res = gp_minimize(
            objective,
            [Real(8, 40)],
            n_calls=50,
            random_state=42,
            acq_func='EI',
            verbose=False
        )
        best_time = round(res.x[0], 1)
        
        groups.append({
            'BMI范围': f"{lower:.1f}-{upper:.1f}",
            '最佳检测孕周': best_time,
            '样本量': len(group_data),
            '预期达标率': round(predict_detection_probability(group_data, best_time), 3),
            '风险等级': '低' if best_time <=12 else ('中' if best_time <=27 else '高')
        })
    
    # 结果展示
    results = pd.DataFrame(groups)
    print("\n=== 最佳分组策略 ===")
    print(results.to_markdown(index=False))
    
    # 可视化
    plot_results(results)
    
    return results

def plot_results(results: pd.DataFrame) -> None:
    """可视化结果"""
    plt.figure(figsize=(14, 7))
    
    # 创建颜色映射和样式
    color_map = {'低': '#4CAF50', '中': '#FFC107', '高': '#F44336'}
    bar_width = 0.6
    
    # 绘制柱状图
    for _, group in results.iterrows():
        plt.bar(
            group['BMI范围'], 
            group['最佳检测孕周'],
            width=bar_width,
            label=f"{group['风险等级']}风险",
            color=color_map[group['风险等级']],
            alpha=0.8,
            edgecolor='black'
        )
        
        # 添加数据标签
        plt.text(
            group['BMI范围'], 
            group['最佳检测孕周'] + 0.3,
            f"{group['最佳检测孕周']}周\n达标率:{group['预期达标率']*100:.1f}%",
            ha='center',
            va='bottom',
            fontsize=10,
            bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=0.5)
        )
    
    # 美化图表
    plt.xlabel('BMI分组', fontsize=12, labelpad=10)
    plt.ylabel('推荐检测孕周', fontsize=12, labelpad=10)
    plt.title('基于BMI的最佳NIPT检测时间推荐 (遗传算法优化)', fontsize=14, pad=20, fontweight='bold')
    
    # 调整图例和网格
    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    plt.legend(by_label.values(), by_label.keys(), title='风险等级', loc='upper right', framealpha=0.9)
    plt.grid(axis='y', linestyle='--', alpha=0.5)
    
    # 调整y轴范围
    y_max = results['最佳检测孕周'].max() + 2
    plt.ylim(8, y_max if y_max > 20 else 20)
    
    plt.tight_layout()
    plt.savefig('ques3_optimal_bmi_groups_ga.png', dpi=300, bbox_inches='tight')
    plt.close()

if __name__ == '__main__':
    df = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')

    main(df)